#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>

#include "FFTEngine.h"

	/** Constructors **/
	FFTEngine::FFTEngine(int iNumPoints, CommLayer *iComm) {
		numPoints = iNumPoints;
		comm = iComm;
		
		rank = comm->rank();
		numNodes = comm->procs();
		numLocalPoints = numPoints / numNodes;
		
		localData = (Complex *) malloc(sizeof(Complex)*numLocalPoints);
		localBuff = (Complex *) malloc(sizeof(Complex)*numLocalPoints);

		if (sizeof(Complex) != sizeof(double) * 2) {
			fprintf(stderr,
				"Error: sizeof(Complex)=%d does not equal 2*sizeof(double)=%d\n",
				(int)sizeof(Complex),(int)sizeof(double)*2);
		}
	}

	FFTEngine::~FFTEngine() {
		free(localData);
		free(localBuff);
	}







	/************************************************ 
		Data Distribution 
	*************************************************/
	
	void FFTEngine::aquireInitData() {
		distributeData(NULL,true,true,true);	
	}

	void FFTEngine::distributeInitData(Complex *arr) {
		distributeData(arr,true,true,true);	
	}

	void FFTEngine::distributeResultData() {
		distributeData(NULL,false,true,false);	
	}

	void FFTEngine::acquireResultData(Complex *results) {
		distributeData(results,false,true,false);
	}







	int FFTEngine::bitRev(int oldVal, int maxVal) {
		int newVal = 0;
		int i;
		int oldBit = maxVal >> 1;
		int newBit = 1;

		while (newBit < maxVal) {

			if (oldVal >= oldBit) {
				newVal += newBit;
				oldVal -= oldBit;
			}
	
			oldBit >>= 1;
			newBit <<= 1;
		}

		return newVal;
	}


	void FFTEngine::local_bitReverse(Complex *arr) {

		int n = numPoints;
		for (int i=0; i < n; i++) {
			int rev = bitRev(i,n);
			if (i < rev) {
				fprintf(stderr,"swapping %d with %d\n",i,rev);	 
				arr[i].swap(&(arr[rev])); 
			}
		}	

	}

	void FFTEngine::local_computeFFT(Complex *arr) {

		//Bit Reverse the array
		local_bitReverse(arr);
		
		int lgn =  log2(numPoints);

		for (int s=1; s <= lgn; s++) {
			fprintf(stderr,"s=%d\n",s);

			int m = pow(2,s);
			Complex wm = FFTEngine::omega(m);

			for (int k=0; k < numPoints; k += m) {

				fprintf(stderr,"k=%d, m=%d\n",k,m);
				Complex w (1.0,0.0);

				for (int j=0; j < (m/2); j++) {
					
					fprintf(stderr,"j=%d w=",j);
					w.printValue();
					fprintf(stderr,"\n");

					int ind0 = k + j;
					int ind1 = k + j + (m/2);

					Complex u = arr[ind0];
					Complex t = w * arr[ind1];
					
					arr[ind0] = u + t;
					arr[ind1] = u - t;

					w = w * wm;
				}
			}
		}
	}
	void FFTEngine::printLocalData() {
		printf("rank=%d my local data\n",rank);
		Complex::printArray(localData,numLocalPoints);
		printf("\n\n");
	}


	void FFTEngine::runFFT(Complex *arr) {
			
		if (rank == 0) {
			distributeInitData(arr);
		} else {
			aquireInitData();
		}

		executeFFT();

		if (rank == 0) {
			acquireResultData(arr);
		} else {
			distributeResultData();
		}
	}

	void FFTEngine::executeFFT(bool commOnly) {
		
		int lgn =  log2(numPoints);

		for (int s=1; s <= lgn; s++) {
			//printf("rank=%d s=%d\n",rank,s);

			int m = pow(2,s);
			Complex wm = FFTEngine::omega(m);

			for (int k=0; k < numPoints; k += m) {

				//printf("k=%d, m=%d\n",k,m);
				Complex w (1.0,0.0);

				for (int j=0; j < (m/2); j++) {
					
					//w.printValue();
					//printf("\n");

					int ind0 = k + j;
					int ind1 = k + j + (m/2);


					//There are 3 senarios that matter
					//  1) ind0 and ind1 are both local
					//  2) ind0 is local and ind1 is dist
					//  2) ind0 is dist and ind0 is dist
					//
					// if ind0 and ind1 are dist, ignore calculation

					if (isLocal(ind0) && isLocal(ind1)) {


						fprintf(stderr,"rank=%d s=%d arr[%d] with arr[%d] LOCALLY\n",rank,s,ind0,ind1);
						int ind0_local = global2localId(ind0);
						int ind1_local = global2localId(ind1);

						//printf("\nind0=%d ind0_local=%d",ind0,ind0_local);
						//printf("  ind1=%d ind1_local=%d\n",ind1,ind1_local);

						Complex u = localData[ind0_local];
						Complex t = w*localData[ind1_local];
					
						localData[ind0_local] = u + t;
						localData[ind1_local] = u - t;
						
					} else if (isLocal(ind0)) {
						Complex u;
						Complex t;

						//Get U
						int ind0_local = global2localId(ind0);
						u = localData[ind0_local];
					
						//Find node to communicate with
						int exNode = globalNode(ind1);
	
						//send U then recieve T
						fprintf(stderr,"rank=%d s=%d arr[%d] with arr[%d] sending U\n",rank,s,ind0,ind1);
						comm->sendComplex(&u,exNode,ind0);
						fprintf(stderr,"rank=%d s=%d arr[%d] with arr[%d] recv T\n",rank,s,ind0,ind1);
						comm->recvComplex(&t,exNode,ind1);
					
	
						//calculate result
						t *= w;
						localData[ind0_local] = u + t;
						
						
					} else if (isLocal(ind1)) {
						Complex u;
						Complex t;

						//Get T
						int ind1_local = global2localId(ind1);
						t = localData[ind1_local];
					
						//Find node to communicate with
						int exNode = globalNode(ind0);
	
						//recieve U then send T
						fprintf(stderr,"rank=%d s=%d arr[%d] with arr[%d] recv U\n",rank,s,ind0,ind1);
						comm->recvComplex(&u,exNode,ind0);
						fprintf(stderr,"rank=%d s=%d arr[%d] with arr[%d] sending T\n",rank,s,ind0,ind1);
						comm->sendComplex(&t,exNode,ind1);
					
	
						//calculate result
						t *= w;
						localData[ind1_local] = u - t;
						

					}	

					w = w * wm;
				}
			}
		}
	}






	/************************************************ 
		Gather and Scatter Constructs 
	*************************************************/

	void FFTEngine::distributeData(Complex *arr,bool isScatter, bool isCyclic, bool isBitReversed) {
		
		//check for sender
		if (rank == 0) {
	
			for (int i=0; i < numPoints; i++) {


				//Caclulate index
				//  when isBitReversed sendIndex != globalIndex
				int sendIndex = i;
				int globalIndex = i;
				
				if (isBitReversed) { sendIndex = bitRev(i,numPoints); }



				//Calculate node
				int node;  //node that value will to sent to
				int localId; //if local, its localData Index
				
				if (isCyclic) {
					node = cyclic_globalNode(globalIndex);
					localId = cyclic_global2localId(globalIndex);
				} else {
					node = blocked_globalNode(globalIndex);
					localId = blocked_global2localId(globalIndex);
				}		


				//if sending to itself
				if (node == 0) {
					if (isScatter) {
						localData[localId] = arr[sendIndex];
					} else {
						arr[sendIndex] = localData[localId];
					}
				} else {
					if (isScatter) {
						//fprintf(stderr,"Rank=%d sending arr[%d] to rank=%d\n",rank,sendIndex,node);
						comm->sendComplex(&(arr[sendIndex]),node,sendIndex);
					} else {
						//fprintf(stderr,"Rank=%d recieving arr[%d] from rank=%d\n",rank,sendIndex,node);
						comm->recvComplex(&(arr[sendIndex]),node,sendIndex);
					}
				}
			}
		} 


		//receiver
		else {

			int startIndex;
			int inc;
			
			if (isCyclic) {
				startIndex = rank;
				inc = numNodes;
			} else {
				startIndex = rank*numLocalPoints;
				inc = 1;
			}
			
			for (int i=0; i < numLocalPoints; i++) {
				int globalId = startIndex + (i*inc);
				int revIndex = globalId;
				int localId;

				//bit reverse
				if (isBitReversed) { revIndex = bitRev(globalId,numPoints); }

				//calculate localId
				if (isCyclic) {
					localId = cyclic_global2localId(globalId);
				} else {
					localId = blocked_global2localId(globalId);
				}
				
				if (isScatter) {
					//fprintf(stderr,"Rank=%d requesting arr[%d]\n",rank,revIndex);
					comm->recvComplex(&(localData[localId]),0,revIndex);
				} else {
					//fprintf(stderr,"Rank=%d sending arr[%d]\n",rank,revIndex);
					comm->sendComplex(&(localData[localId]),0,revIndex);
				}
			}
		}

	}













	/************************************************ 
		Utility Functions 
	*************************************************/

	void FFTEngine::revCopy(Complex *arr, int n) {

		for (int i=0; i < n; i++) {
			int revIndex = bitRev(i,n);
			if (i < revIndex) {
				//printf("swapping %d with %d\n",i,rev(i,n));	 
				arr[i].swap(&(arr[revIndex])); 
			}
		}
	}


	int FFTEngine::pow(int v, int p) {
		int prod = 1;
		for (int i=0; i < p; i++) {
			prod *= v;
		}
		return prod;
	}

	int FFTEngine::log2(int v) {
		int count = 0;
		int val = 1;
		
		while (v > val) {
			val *= 2;
			count++;
		}

		return count;
	}



	/************************************************ 
		Unit Tests
	*************************************************/
	void FFTEngine::unitTests() {
		testDistributeData();
		testFFT();
	}

	void complexAssert(Complex a, Complex b, int rank, int index) {

		if (!(a == b)) {
			printf("rank=%d checking arr[%d] ",rank,index);
			a.printValue();
			printf(" != ");
			b.printValue();
			printf("\n");
		}
		assert(a == b);
	}

	void FFTEngine::testFFT() {
	
		//Create data structures
		Complex *startVals = (Complex*) malloc(numPoints * sizeof(Complex));
		Complex *localResults = (Complex*) malloc(numPoints * sizeof(Complex));
		Complex *distResults = (Complex*) malloc(numPoints * sizeof(Complex));


		//create values
		for (int i=0; i < numPoints; i++) { startVals[i] = Complex((i+1)*1.0,0.0); }
		for (int i=0; i < numPoints; i++) { localResults[i] = Complex((i+1)*1.0,0.0); }


		
		if (rank == 0) {
			distributeInitData(startVals);
		} else {
			aquireInitData();
		}

		executeFFT();

		if (rank == 0) {
			acquireResultData(distResults);
			local_computeFFT(localResults);

			for (int i=0; i < numPoints; i++) {
				complexAssert(distResults[i],localResults[i],rank,i);
			}

		} else {
			distributeResultData();
		}

		//
		// All tests passed!
		//
		comm->barrier(); //Wait before declaring success
		if (rank==0) { printf("\n\ntestFFT() passed\n"); }


		free(startVals);
		free(localResults);
		free(distResults);
	}

	void FFTEngine::testDistributeData() {
		
		// Testing
		//   distributeData(Complex *startVals,
		//			bool isScatter, 
		//			bool isCyclic, 
		//			bool isBitReversed);



		//Create data structures
		Complex *startVals = (Complex*) malloc(numPoints * sizeof(Complex));
		Complex *endVals = (Complex*) malloc(numPoints * sizeof(Complex));


		//
		//test cyclic scatter
		//
		for (int i=0; i < numPoints; i++) { startVals[i] = Complex((i+1)*1.0,0.0); }
		distributeData(startVals,true,true,false);
		
		for (int i=0; i < numPoints; i++) {
			if (cyclic_isLocal(i)) {
				int localIndex = cyclic_global2localId(i);
				complexAssert(startVals[i],localData[localIndex],rank,i);
			}
		}

		comm->barrier(); //Wait before declaring success
		if (rank==0) { printf("...cyclic scatter passed\n"); }



		//test cyclic gather
		for (int i=0; i < numPoints; i++) { endVals[i] = Complex(-1.5,-1.5); }
		distributeData(endVals,false,true,false);

		for (int i=0; i < numPoints; i++) {
			if (rank == 0) {
				complexAssert(startVals[i],endVals[i],rank,i);
			}
		}

		comm->barrier(); //Wait before declaring success
		if (rank==0) { printf("...cyclic gather passed\n"); }







		//
		//test cyclic scatter with bit reverse
		//
		for (int i=0; i < numPoints; i++) { startVals[i] = Complex((i+1)*1.0,0.0); }
		distributeData(startVals,true,true,true);
		local_bitReverse(startVals);
		
		for (int i=0; i < numPoints; i++) {
			if (cyclic_isLocal(i)) {
				int localIndex = cyclic_global2localId(i);
				complexAssert(startVals[i],localData[localIndex],rank,i);
			}
		}

		comm->barrier(); //Wait before declaring success
		if (rank==0) { printf("...cyclic scatter with bit reverse passed\n"); }




		//test cyclic gather with bit reverse
		for (int i=0; i < numPoints; i++) { endVals[i] = Complex(-1.5,-1.5); }
		distributeData(endVals,false,true,true);

		for (int i=0; i < numPoints; i++) {
			if (rank == 0) {
				complexAssert(startVals[bitRev(i,numPoints)],endVals[i],rank,i);
			}
		}

		comm->barrier(); //Wait before declaring success
		if (rank==0) { printf("...cyclic gather with bit reverse passed\n"); }





		//
		//test blocked scatter
		//
		for (int i=0; i < numPoints; i++) { startVals[i] = Complex((i+1)*1.0,0.0); }
		distributeData(startVals,true,false,false);
		
		for (int i=0; i < numPoints; i++) {
			if (blocked_isLocal(i)) {
				int localIndex = blocked_global2localId(i);
				complexAssert(startVals[i],localData[localIndex],rank,i);
			}
		}

		comm->barrier(); //Wait before declaring success
		if (rank==0) { printf("...blocked scatter passed\n"); }


		//test cyclic gather
		for (int i=0; i < numPoints; i++) { endVals[i] = Complex(-1.5,-1.5); }
		distributeData(endVals,false,false,false);

		for (int i=0; i < numPoints; i++) {
			if (rank == 0) {
				complexAssert(startVals[i],endVals[i],rank,i);
			}
		}

		comm->barrier(); //Wait before declaring success
		if (rank==0) { printf("...blocked gather passed\n"); }





		//
		//test blocked scatter with bit reverse
		//
		for (int i=0; i < numPoints; i++) { startVals[i] = Complex((i+1)*1.0,0.0); }
		distributeData(startVals,true,false,true);
		local_bitReverse(startVals);
		
		for (int i=0; i < numPoints; i++) {
			if (blocked_isLocal(i)) {
				int localIndex = blocked_global2localId(i);
				complexAssert(startVals[i],localData[localIndex],rank,i);
			}
		}

		comm->barrier(); //Wait before declaring success
		if (rank==0) { printf("...blocked scatter with bit reverse passed\n"); }


		//test cyclic gather
		for (int i=0; i < numPoints; i++) { endVals[i] = Complex(-1.5,-1.5); }
		distributeData(endVals,false,false,true);

		for (int i=0; i < numPoints; i++) {
			if (rank == 0) {
				complexAssert(startVals[bitRev(i,numPoints)],endVals[i],rank,i);
			}
		}

		comm->barrier(); //Wait before declaring success
		if (rank==0) { printf("...blocked gather with bit reverse passed\n"); }





		//
		// All tests passed!
		//
		comm->barrier(); //Wait before declaring success
		if (rank==0) { printf("\n\ntestDistributeData() passed\n"); }


		free(startVals);
		free(endVals);
	}

